from datetime import datetime
from collections import deque
from typing import Dict, Any
from rainforeLearn.gomoku.v2.train.constants.train_constants import TrainConstants


class TrainingStateManager:
    """训练状态管理器 - 负责训练状态的保存和恢复"""
    
    def __init__(self, trainer):
        self.trainer = trainer
    
    def get_current_state(self) -> Dict[str, Any]:
        """获取当前训练状态"""
        elapsed_time = self._calculate_elapsed_time()
        
        return {
            'current_episode': self.trainer.current_episode,
            'total_episodes': self.trainer.total_episodes,
            'episode_rewards': list(self.trainer.episode_rewards),
            'episode_lengths': list(self.trainer.episode_lengths),
            'win_rates': list(self.trainer.win_rates),
            'best_win_rate': self.trainer.best_win_rate,
            'best_episode': self.trainer.best_episode,
            'total_training_time': self.trainer.total_training_time + elapsed_time,
            'config': self.trainer.config.__dict__,
            'timestamp': datetime.now().isoformat(),
            'reward_system_state': self._get_reward_system_state(),
            'stats_printer_state': self._get_stats_printer_state()
        }
    
    def restore_state(self, trainer_state: Dict[str, Any]) -> None:
        """恢复训练状态"""
        self._restore_basic_state(trainer_state)
        self._restore_reward_system_state(trainer_state.get('reward_system_state', {}))
        self._restore_stats_printer_state(trainer_state.get('stats_printer_state', {}))
    
    def _calculate_elapsed_time(self) -> float:
        """计算已用时间"""
        if self.trainer.training_start_time:
            return (datetime.now() - self.trainer.training_start_time).total_seconds()
        return 0.0
    
    def _restore_basic_state(self, trainer_state: Dict[str, Any]) -> None:
        """恢复基础状态"""
        self.trainer.current_episode = trainer_state['current_episode']
        self.trainer.total_episodes = trainer_state['total_episodes']
        self.trainer.episode_rewards = deque(trainer_state['episode_rewards'], maxlen=TrainConstants.DEFAULT_BUFFER_SIZE)
        self.trainer.episode_lengths = deque(trainer_state['episode_lengths'], maxlen=TrainConstants.DEFAULT_BUFFER_SIZE)
        self.trainer.win_rates = deque(trainer_state['win_rates'], maxlen=TrainConstants.DEFAULT_WIN_RATES_SIZE)
        self.trainer.best_win_rate = trainer_state['best_win_rate']
        self.trainer.best_episode = trainer_state['best_episode']
        self.trainer.total_training_time = trainer_state.get('total_training_time', 0.0)
    
    def _get_reward_system_state(self) -> Dict[str, Any]:
        """获取奖励系统状态"""
        reward_system = self.trainer.reward_system
        return {
            'batch_win_rates': reward_system.batch_stats.win_rates,
            'batch_games_results': reward_system.batch_stats.game_results,
            'win_rate_bonus_awarded': reward_system.win_rate_bonus_awarded,
            'batch_penalty_history': reward_system.batch_penalty_history,
            'penalty_history': reward_system.penalty_history,
            'total_batch_bonus': reward_system.total_batch_bonus,
            'total_batch_penalties': reward_system.total_batch_penalties
        }
    
    def _restore_reward_system_state(self, reward_state: Dict[str, Any]) -> None:
        """恢复奖励系统状态"""
        reward_system = self.trainer.reward_system
        reward_system.batch_stats.win_rates = reward_state.get('batch_win_rates', [])
        reward_system.batch_stats.game_results = reward_state.get('batch_games_results', [])
        reward_system.win_rate_bonus_awarded = reward_state.get('win_rate_bonus_awarded', [])
        reward_system.batch_penalty_history = reward_state.get('batch_penalty_history', [])
        reward_system.penalty_history = reward_state.get('penalty_history', [])
        reward_system.batch_stats.comparisons = reward_state.get('recent_batch_comparisons', [])
        reward_system.total_batch_bonus = reward_state.get('total_batch_bonus', 0.0)
        reward_system.total_batch_penalties = reward_state.get('total_batch_penalties', 0.0)

    def _get_stats_printer_state(self) -> Dict[str, Any]:
        """获取统计打印器状态"""
        stats_printer = self.trainer.stats_printer
        return {
            'total_wins': stats_printer.game_tracker.total_wins,
            'total_losses': stats_printer.game_tracker.total_losses,
            'total_draws': stats_printer.game_tracker.total_draws,
            'win_rate_history': stats_printer.game_tracker.win_rate_history,
            'total_batch_bonus': stats_printer.reward_penalty_tracker.total_batch_bonus,
            'win_rate_bonus_history': stats_printer.reward_penalty_tracker.win_rate_bonus_history,
            'total_batch_penalties': stats_printer.reward_penalty_tracker.total_batch_penalties,
            'batch_penalty_history': stats_printer.reward_penalty_tracker.batch_penalty_history,
            'threat_defense_stats': stats_printer.threat_tracker.to_dict(),
            'invalid_moves_per_episode': stats_printer.game_tracker.invalid_moves_per_episode
        }

    def _restore_stats_printer_state(self, stats_state: Dict[str, Any]) -> None:
        """恢复统计打印器状态"""
        stats_printer = self.trainer.stats_printer
        stats_printer.game_tracker.total_wins = stats_state.get('total_wins', 0)
        stats_printer.game_tracker.total_losses = stats_state.get('total_losses', 0)
        stats_printer.game_tracker.total_draws = stats_state.get('total_draws', 0)
        stats_printer.game_tracker.win_rate_history = stats_state.get('win_rate_history', [])
        stats_printer.reward_penalty_tracker.total_batch_bonus = stats_state.get('total_batch_bonus', 0)
        stats_printer.reward_penalty_tracker.win_rate_bonus_history = stats_state.get('win_rate_bonus_history', [])
        stats_printer.reward_penalty_tracker.total_batch_penalties = stats_state.get('total_batch_penalties', 0)
        stats_printer.reward_penalty_tracker.batch_penalty_history = stats_state.get('batch_penalty_history', [])
        # threat_defense_stats 不需要恢复，因为它是通过to_dict()动态生成的
        stats_printer.game_tracker.invalid_moves_per_episode = stats_state.get('invalid_moves_per_episode', [])